# -*- coding: UTF-8 -*-
import struct
#from bp import *
from datetime import datetime

import gzip
import os
import urllib

import numpy

g_rows = 56
g_cols = 28

# 数据加载器基类
class Loader(object):

    def __init__(self, path, count):
        '''
        初始化加载器
        path: 数据文件路径
        count: 文件中的样本个数
        '''
        self.path = path
        self.count = count

    def __init__(self,path):
        self.path = path

    def get_file_content(self):
        '''
        读取文件内容
        '''
        f = open(self.path, 'rb')
        content = f.read()
        f.close()
        return content

    def to_int(self, byte):
        '''
        将unsigned byte字符转换为整数
        '''
        return struct.unpack('B', byte)[0]


def _read32(bytestream):
  dt = numpy.dtype(numpy.uint32).newbyteorder('>')
  #return numpy.frombuffer(bytestream.read(4), dtype=dt) old ver.
  return numpy.frombuffer(bytestream.read(4), dtype=dt)[0]


# 图像数据加载器
class ImageLoader(Loader):

    def extract_images(self,filename):
        """Extract the images into a 4D uint8 numpy array [index, y, x, depth]."""
        print('Extracting', filename)

        image_bin_file = open(filename, 'rb')

        buffer = image_bin_file.read(16)

        magic, number, rows, cols = struct.unpack('iiii', buffer)

        image_bin_file.seek(16)

        buffer2 = image_bin_file.read(number * rows * cols)

        dat = struct.unpack(str(number * rows * cols) + 'B', buffer2)

        image_dat = numpy.reshape(dat, [number, rows * cols])

        image_bin_file.close()

        return image_dat

    def get_picture(self, content, index):
        '''
        内部函数，从文件中获取图像
        '''
        start = index * g_rows * g_cols + 16
        picture = []
        for i in range(g_rows):
            picture.append([])
            for j in range(g_cols):
                picture[i].append(
                    self.to_int(content[start + i * g_rows + j]))
        return picture

    def get_one_sample(self, picture):
        '''
        内部函数，将图像转化为样本的输入向量
        '''
        sample = []
        for i in range(g_rows):
            for j in range(g_cols):
                sample.append(picture[i][j])
        return sample

    def load(self):
        '''
        加载数据文件，获得全部样本的输入向量
        '''
        content = self.get_file_content()
        data_set = []
        for index in range(self.count):
            data_set.append(
                self.get_one_sample(
                    self.get_picture(content, index)))
        return data_set

    def loadx(self):
        return self.extract_images(self.path)


# 标签数据加载器
class LabelLoader(Loader):

    def dense_to_one_hot(self, labels_dense, num_classes):
        """Convert class labels from scalars to one-hot vectors."""
        num_labels = labels_dense.shape[0]
        index_offset = numpy.arange(num_labels) * num_classes
        labels_one_hot = numpy.zeros((num_labels, num_classes))
        labels_one_hot.flat[index_offset + labels_dense.ravel()] = 1
        return labels_one_hot

    def extract_labels(self, filename, one_hot=False,n_classes = 10):
        """Extract the labels into a 1D uint8 numpy array [index]."""
        print('Extracting', filename)

        label_bin_file = open(filename, "rb")

        buffer = label_bin_file.read(8)

        magic, number = struct.unpack('ii', buffer)

        label_bin_file.seek(8)

        buffer2 = label_bin_file.read(number)

        labels = numpy.reshape(struct.unpack(str(number) + 'B', buffer2),-1)

        label_bin_file.close()

        if one_hot:
            return self.dense_to_one_hot(labels,n_classes)

        return labels

    def load(self):
        '''
        加载数据文件，获得全部样本的标签向量
        '''
        content = self.get_file_content()
        labels = []
        for index in range(self.count):
            labels.append(self.norm(content[index + 8]))
        return labels

    def loadx(self,one_hot,n_classes):
        '''
        改写一下load
        :return:
        '''
        return self.extract_labels(self.path,one_hot,n_classes)

    def norm(self, label):
        '''
        内部函数，将一个值转换为10维标签向量
        '''
        label_vec = []
        label_value = self.to_int(label)
        for i in range(10):
            if i == label_value:
                label_vec.append(0.9)
            else:
                label_vec.append(0.1)
        return label_vec


datafolder = '/Users/vista/PycharmProjects/data/ocr_data'

images_file_name_digit = 'digit-images-ubyte'
labels_file_name_digit = 'digit-labels-ubyte'
images_file_name_letter = 'letter-images-ubyte'
labels_file_name_letter = 'letter-labels-ubyte'

def get_training_data_set():
    '''
    获得训练数据集
    '''
    image_loader = ImageLoader( datafolder + 'train-images-idx3-ubyte.gz', 60000)
    label_loader = LabelLoader( datafolder + 'train-labels-idx1-ubyte.gz', 60000)
    return image_loader.load(), label_loader.load()


def get_training_data_set_x(one_hot = False):
    '''
    获得训练数据集
    '''
    image_loader = ImageLoader( datafolder + 'train-images-idx3-ubyte.gz')
    label_loader = LabelLoader( datafolder + 'train-labels-idx1-ubyte.gz')
    return image_loader.loadx(), label_loader.loadx(one_hot)

import data.dataset as dd

def get_mnist_set(one_hot = False):
    '''
    获得训练数据集
    '''
    trimages_d = ImageLoader( os.path.join(datafolder, images_file_name_digit) ).loadx()
    trlabels_d = LabelLoader( os.path.join(datafolder, labels_file_name_digit) ).loadx(one_hot,10)

    trimages_l = ImageLoader( os.path.join(datafolder, images_file_name_letter) ).loadx()
    trlabels_l = LabelLoader( os.path.join(datafolder, labels_file_name_letter) ).loadx(one_hot,26)

    # read train data
    class DataSets(object):
        pass
    data_sets = DataSets()

    data_sets.digit = dd.DataSet(trimages_d,trlabels_d)
    data_sets.letter = dd.DataSet(trimages_l, trlabels_l)

    return data_sets